from PIL import Image, ImageDraw
import os
from tqdm import tqdm
import xml.etree.ElementTree as ET

# Set image-input path
folder_path = 'F:\\items\\bike-model\\搜集的数据集\\自行制作的\\zzxi'

# Set annotation-input path
annotation_input_path = 'F:\\items\\bike-model\\搜集的数据集\\自行制作的\\zzxa'

# set painted-image path
painted_img_path = 'F:\\items\\bike-model\\搜集的数据集\\自行制作的\\zzxs'

for filename in tqdm(os.listdir(folder_path)):
    if filename.split('.')[-1] == 'jpg':
        # Load the image
        file_path = os.path.join(folder_path, filename)
        # Load the image
        image = Image.open(file_path)

        # Create a drawing object
        draw = ImageDraw.Draw(image)

        # Load the VOC XML file
        tree = ET.parse(
            os.path.join(annotation_input_path, filename.split('.')[0] + '.xml'))
        root = tree.getroot()

        # Draw the bounding boxes
        for obj in root.findall('object'):
            bbox = obj.find('bndbox')
            xmin = int(bbox.find('xmin').text)
            ymin = int(bbox.find('ymin').text)
            xmax = int(bbox.find('xmax').text)
            ymax = int(bbox.find('ymax').text)
            draw.rectangle([xmin, ymin, xmax, ymax], outline='red')

        # Save the annotated image
        image.save(os.path.join(painted_img_path, filename))
